# -*- coding: utf-8 -*-
from gwac_common import *
import numpy as np
import os
import sys
import glob
import requests
import shutil
import traceback
from datetime import datetime 
iptable=[
    'http://190.168.1.11/split/',
    'http://190.168.1.12/split/',
    'http://190.168.1.13/split/',
    ]

def zonefalpha(theta, fdec):
    if np.abs(fdec)+theta>89.9:return 180.0;
    return np.rad2deg(np.abs(np.arctan(np.sin(np.deg2rad(theta))/ np.sqrt(np.abs(np.cos(np.deg2rad(fdec-theta)) * np.cos(np.deg2rad(fdec+theta)))))))

def zonefalpha2(theta, fdec):
    return np.rad2deg(np.abs(np.arctan(np.sin(np.deg2rad(theta))/ np.sqrt(np.abs(np.cos(np.deg2rad(fdec-theta)) * np.cos(np.deg2rad(fdec+theta)))))))

#处理zoneheihgt远小于theta的情形，即把theta当作zoneheight，而当前fdec是该zone的最靠近极点的一边
#——纯粹是想多了
def zonefalphax(theta, fdec):
    ht=theta/2.0
    if fdec>ht:
        fdec-=theta
    if fdec<-ht:
        fdec+=theta
    return zonefalpha(theta, fdec)

def testZoneAlpha(theta=10):
    print('theta=',theta)
    for d in range(-92, 93):
        print(d, zonefalpha2(theta, d), zonefalpha(theta, d))

def subPath2Param(subpath):
    dec,rest=subpath.split('/')
    ra,rest=rest.split('_')
    mag,rest=rest.split('.')
    return int(dec),int(ra),int(mag)

def rectangle(dec, ra, fieldwidth):
    hfw=fieldwidth/2.0
    alpha=zonefalpha(hfw, dec)
    dec_min=dec-hfw
    if dec_min <-90: dec_min=-90
    dec_max=dec+hfw
    if dec_max >=90: dec_max=89.9
    ra_min=ra-alpha
    ra_max=ra+alpha
    if ra_min<0:
        return [
            [dec_min, dec_max, ra_min+360, 360],
            [dec_min, dec_max, 0, ra_max]
            ]
    elif ra_max>360:
        return [
            [dec_min, dec_max, ra_min, 360],
            [dec_min, dec_max, 0, ra_max-360]
            ]
    return [
        [dec_min, dec_max, ra_min, ra_max]
        ]

def blockList(dec, ra, fieldwidth):
    for dec_min, dec_max, ra_min, ra_max in rectangle(dec, ra, fieldwidth):
        decrange=range(int(np.floor(dec_min)), int(np.floor(dec_max))+1)
        rarange=range(int(np.floor(ra_min)), int(np.floor(ra_max))+1)
        for dec in decrange:
            for ra in rarange:
                yield (dec, ra)

#dec应该是一个区块靠近南/北极点的赤纬
def calcStripeRARange(dec, ra, width):
    alpha=zonefalpha(width, dec)
    ra_min=ra-alpha
    ra_max=ra+alpha
    if ra_min<0:
        return [
            [ra_min+360, 360],
            [0, ra_max]
            ]
    elif ra_max>360:
        return [
            [ra_min, 360],
            [0, ra_max-360]
            ]
    return [
        [ra_min, ra_max]
        ]

def stripeBlocksList(cdec, cra, fieldwidth):
    hfw=fieldwidth/2.0
    lowdec,highdec=int(np.floor(cdec-hfw)), int(np.floor(cdec+hfw)+1)
    if lowdec<-90:
        lowdec=-90
    if highdec>90:
        highdec=90
    for dec in range(lowdec, highdec):
        for ra_min, ra_max in calcStripeRARange(dec, cra, hfw):
            for ra in range(int(np.floor(ra_min)), int(np.floor(ra_max))+1):
                yield (dec, ra)

def circleRARange(dec, ra, dec1, radius):
    if dec1>90:
        dec1=180-dec1
    elif dec1<-90:
        dec1=-180-dec1
    ra=np.deg2rad(ra)
    dec=np.deg2rad(dec)
    dec1=np.deg2rad(dec1)
    radius=np.deg2rad(radius)
    #
    #h=1-np.sin(dec)*np.sin(dec1)-2*np.sin(radius/2.0)**2
    h=np.cos(radius)-np.sin(dec)*np.sin(dec1)
    #print('h=',h)
    h=h/np.cos(dec)/np.cos(dec1)
    err=False
    if h<-1:
        if h>-1.00000000000001:
            h=-1
        else:
            err=True
    if h>1:
        if h< 1.00000000000001:
            h=1
        else:
            err=True
    if err:
        #print('Error dec1={0}, h={1}'.format(np.rad2deg(dec1), h))
        return [[0,360]]
    ra1=np.arccos(h)
    ra_min,ra_max=np.rad2deg(ra-ra1),np.rad2deg(ra+ra1)
    if ra_min<0:
        return [[ra_min+360, 360], [0, ra_max]]
    elif ra_max>360:
        return [[ra_min, 360],[0, ra_max-360]]
    else: return [[ra_min, ra_max]]

def circleSearch(dec, ra, radius):
    dec_min=int(np.floor(dec-radius))
    dec_max=int(np.floor(dec+radius))
    #print('dec_min={0},dec_max={1}'.format(dec_min,dec_max))
    if dec_min<-90:
        cover=int(np.floor(-180-dec+radius))
        for d in range(-90, cover+1):
            for r in range(0,360):
                yield d,r
        dec_min=cover+1
    if dec_max>90:
        cover=int(np.floor(180-dec-radius))
        for d in range(cover, 90):
            for r in range(0,360):
                yield d,r
        dec_max=cover-1
    #print('dec_min={0},dec_max={1}'.format(dec_min,dec_max))
    for d in range(dec_min, dec_max+1):
        if d>dec:
            real_d=d-1
        elif d<dec:
            real_d=d+1
        else:
            real_d=d
        for ra_min,ra_max in circleRARange(dec, ra, real_d, radius):
            #print(d,ra_min,ra_max)
            ra_min=int(np.floor(ra_min))
            ra_max=int(np.floor(ra_max))
            for r in range(ra_min,ra_max + 1):
                yield d,r

def download(url, filename):
    makeParentDirs(filename)
    try:
        r = requests.get(url,data={}, timeout=240)
        with open(filename, 'wb') as f:
            for chunk in r.iter_content(chunk_size=1024):
                if chunk: # filter out keep-alive new chunks
                    f.write(chunk)
    except:
        print("Error when handling with url=",url,"filename=",filename)
        print("More Error Info:", traceback.extract_stack())

def copyFile(srcpath, destpath):
    if not os.path.exists(srcpath):
        #print('file not exists', srcpath)
        return
    makeParentDirs(destpath)
    with open(destpath, 'wb') as wfd:
        with open(srcpath, 'rb') as fd:
            shutil.copyfileobj(fd, wfd, 1024*1024*10)
        #wfd.flush()

def getFileList(dirpath, pattern='*.csv'):
    filelist=glob.glob(dirpath+pattern)
    return filelist

def mergeCsvs(destfile, filelist,header):
    makeParentDirs(destfile)
    with open(destfile, 'wb') as wfd:
        if header is not None:
            wfd.write(header)
        for rf in filelist:
            if not os.path.exists(rf):
                continue
            with open(rf, 'rb') as fd:
                #10MB per writing chunk to avoid reading big file into memory.
                shutil.copyfileobj(fd, wfd, 1024*1024*10)
        #wfd.flush()

def substractSet(a, b):
    if not isinstance(a, set):
        a=set(a)
    if not isinstance(b, set):
        b=set(b)
    return a-b

def travaseDir(rootdir, ext='.csv'):
    files = [ 
           os.path.join(parent, name)
           for (parent, subdirs, files) in os.walk(rootdir)
           for name in files + subdirs
       ]
    files = filter(lambda x:x.endswith(ext), files)
    return files

def rectSearch(ra,dec,fieldwidth,mag,mergedFileName='merge.csv'):
    start_time = datetime.now()
    print('{0} Start'.format(start_time))
    ####
    #header=b'source_id,healpixid,htmid,ra,dec,pmra,pmdec,parallax,g_mag,bp_mag,rp_mag,teff,ra2015,dec2015,parallax2015,pmra2015,pmdec2015\n'
    header=b'source_id,healpixid,htmid,ra,dec,pmra,pmdec,parallax,g_mag,bp_mag,rp_mag,teff,ra2015,dec2015,parallax2015,parallax_error2015,phot_g_mean_flux2015,phot_g_mean_flux_error2015,pmra2015,pmdec2015\n'
    #cachepath='/data/gaia_dr2/cached-catalog/'
    cachepath='/home/w60ccd/workspace/gaia_dr2/cached-catalog/'
    datasetName='gaia_dr2'
    mergedFileName='{0}{1}/{2}'.format(cachepath,datasetName,mergedFileName)
    #ra,dec=230,30
    #60cm
    #fieldwidth,mag=19.0/60.0, 21
    #30cm
    #fieldwidth,mag=2, 21
    #JFoV
    #fieldwidth,mag=13, 19
    #FFoV
    #fieldwidth,mag=26, 16
    #print('ra={0}, dec={1}, fieldwidth={2}, mag={3}'.format(ra, dec, fieldwidth, mag))
    mag-=0.01
    #
    if   mag<13: magrange=[0]
    elif mag<16: magrange=[0,1]
    elif mag<19: magrange=[0,1,2]
    else: magrange=[0,1,2,3]
    #
    #existingFileList=glob.glob(cachepath+'**/*.csv', recursive=True)
    existingFileList=travaseDir(cachepath,'.csv')
    for i in range(0, len(existingFileList)):
        existingFileList[i]=existingFileList[i].replace('\\','/').replace('//','/')
    #
    machineNumber=len(iptable)
    expectFileList=[]
    for d,r in circleSearch(dec, ra, fieldwidth/2.0):
    #for d, r in stripeBlocksList(dec, ra, fieldwidth):
    #for d, r in blockList(dec, ra, fieldwidth):
        for magid in magrange:
            expect="{0}{1}/{2}/{3}_{4}.csv".format(cachepath,datasetName,d,r,magid)
            #print(expect)
            expectFileList.append(expect)
            if expect in existingFileList:
                continue
            machine=int((d+r+magid)%machineNumber)
            ip=iptable[machine]
            url='{0}{1}/{2}/{3}_{4}.csv'.format(ip, datasetName,d, r, magid)
            if url.startswith('D:'):
                print('copy', expect, 'from machine', machine, url)
                copyFile(url, expect)
            else:
                #print('download', expect, 'from', url)
                download(url, expect)
    #
    '''
    #extra files多余的文件
    extraFiles=substractSet(existingFileList, expectFileList)
    for f in extraFiles:
        if f.find('merge-')>0:
            continue
        if os.path.exists(f):
            #print('removing', f)
            #os.remove(f)
            pass
    '''
    #merge csvs合并文件
    print('{0} merging files'.format(datetime.now()))
    mergeCsvs(mergedFileName, expectFileList, header)
    end_time = datetime.now()
    print('{0} End'.format(end_time))
    print('Done! Time elapsed {0}'.format(end_time-start_time))
    print('########################')

####
#testZoneAlpha(theta=10)
#ra,dec=230,30
#ra,dec=300,20
#ra,dec=120.13,76.5
#ra,dec=45,8.5
fieldwidth,mag,mergedFileName=26,13,'merge-ffov.csv'
#fieldwidth,mag,mergedFileName=13, 19, 'merge-jfov.csv'
#fieldwidth,mag,mergedFileName=19.0/60.0, 21, 'merge-60cm.csv'
#fieldwidth,mag,mergedFileName=2, 21, 'merge-30cm.csv'
argmap={}
for ar in sys.argv:
    ss = ar.split('=')
    if len(ss) != 2:
        continue
    argmap[ ss[0].strip() ]=ss[1].strip('\n').strip('\r').strip(' ').strip('\'')
#
if 'ra'             in argmap: ra              = float(argmap['ra'])
if 'dec'            in argmap: dec             = float(argmap['dec'])
if 'fieldwidth'     in argmap: fieldwidth      = float(argmap['fieldwidth'])
if 'mag'            in argmap: mag             = float(argmap['mag'])
if 'mergedFileName' in argmap: mergedFileName  = argmap['mergedFileName']
#
print('ra:',ra)
print('dec:',dec)
print('fieldwidth:',fieldwidth)
print('mag:',mag)
print('mergedFileName:',mergedFileName)
#
rectSearch(ra,dec,fieldwidth,mag,mergedFileName)
